import numpy as np
import rasterio
from rasterio.windows import Window
import cv2
from shapely.geometry import Polygon, Point
from rasterio.enums import Resampling
import os
def find_corners1(img):
    '''

    :param image:
    :return: corners(左上，右下，列行)
    '''

        # 转换为灰度图
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 应用阈值处理
    _, thresh = cv2.threshold(gray, 2, 253, cv2.THRESH_BINARY)

    # 寻找轮廓
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 假设最大的轮廓是主体内容
    if contours:
        max_contour = max(contours, key=cv2.contourArea)

        # 轮廓的边界点（近似为角点）
        # 这里我们直接使用轮廓的边界点，但在实际情况中可能需要更复杂的算法来定义角点
        peri = cv2.arcLength(max_contour, True)
        approx = cv2.approxPolyDP(max_contour, 0.04 * peri, True)

        # 角点坐标(左上，左下，右下，右上)列行
        corners = approx.reshape((-1, 2))

        return corners


def raster2opencv(image, bands=None):
    """
    读取 TIFF 图像并转换为 OpenCV 可以处理的格式。

    参数:
    file_path (str): TIFF 文件的路径。
    bands (list or None): 需要读取的波段编号，例如 [1, 2, 3] 表示读取第 1、2、3 波段；
                          如果为 None，则默认读取所有波段。

    返回:
    image (np.ndarray): 转换后的图像，形状为 (height, width, channels)，通道顺序为 BGR。
    """
    # with rasterio.open(file_path) as src:
    #     # 读取指定波段
    #     if bands is None:
    #         # 读取所有波段
    #         image = src.read()
    #     else:
    #         # 读取指定波段
    #         image = np.stack([src.read(band) for band in bands], axis=0)

    # 调整数据类型
    if image.dtype == 'uint8':
        # 数据已经是 uint8 类型
        pass
    elif image.dtype == 'float32' or image.dtype == 'float64':
        # 将浮点数转换为 uint8 类型
        image = ((image - image.min()) / (image.max() - image.min()) * 255).astype(np.uint8)
    else:
        # 其他数据类型，可以根据需要转换
        image = image.astype(np.uint8)

    # 调整通道顺序为 BGR
    image = np.moveaxis(image, 0, -1)

    # OpenCV 使用 BGR 通道顺序
    if image.shape[-1] == 3:
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

    return image


def save_blocks_based_on_mask(src,block_size,blocks, mask, output_dir, prefix='block_', suffix='.tif', resampling=Resampling.bilinear):
    """
    根据给定的掩码保存分割后的图像块。

    参数:
    blocks (np.ndarray): 分割后的图像块数组，形状为 (num_blocks_y, num_blocks_x)，元素为 NumPy 数组。
    mask (np.ndarray): 掩码数组，形状与 blocks 相同，值为 0 或 1，指示是否保存该块。
    output_dir (str): 输出目录的路径。
    prefix (str): 输出文件名的前缀，默认为 'block_'。
    suffix (str): 输出文件名的后缀，默认为 '.tif'。
    resampling (Resampling): 重采样方法，默认为 Resampling.nearest。

    返回:
    None
    """
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    # with rasterio.open(file_path) as src:
    crs = src.crs
    transform = src.transform
    # 遍历每个块
    for y in range(blocks.shape[0]):
        for x in range(blocks.shape[1]):
            if mask[y, x]:
                # 获取块和保存路径
                block = blocks[y, x]
                start_x = x * block_size
                end_x = min(start_x + block_size, src.width)
                start_y = y * block_size
                end_y = min(start_y + block_size, src.height)
                # 构造文件名
                filename = f"{prefix}_{start_x}_{start_y}_{end_x}_{end_y}{suffix}"
                # out_path = os.path.join(output_dir, filename)
                out_path = os.path.join(output_dir,filename )
                new_transform = transform * transform.translation(start_x, start_y)
                # 创建输出文件
                with rasterio.open(
                        out_path,
                        'w',
                        driver='GTiff',
                        height=block.shape[1],
                        width=block.shape[2],
                        count=src.count,  # 单波段
                        dtype=block.dtype,
                        crs=crs,  # 保持与原图相同的坐标系统
                        transform=new_transform,
                        nodata=None
                ) as dst:
                    dst.write(block)
def read_tif_and_split_into_blocks(file_path, block_size,out_path):
    """
    读取 TIFF 格式的卫星影像，并将其分割为 nxn 的正方形块。

    参数:
    file_path (str): TIFF 文件的路径。
    block_size (int): 正方形块的边长。
    bounding_box (list): 方框 A 的坐标 [x1, y1, x2, y2, x3, y3, x4, y4]。

    返回:
    blocks (np.ndarray): 分割后的正方形块数组。
    centers_in_box (np.ndarray): 判断中心点是否位于方框 A 内的结果数组（值为 0 或 1）。
    blocks_in_box (np.ndarray): 判断正方形块是否完全位于方框 A 内的结果数组（值为 0 或 1）。
    """
    # 读取 TIFF 文件
    with rasterio.open(file_path) as src:
        image = src.read()  # 读取
        _,height,width= image.shape

        bounding_box = find_corners1(raster2opencv(image)).reshape(-1,).tolist()
        # 创建平行四边形的 Shapely Polygon 对象
        poly = Polygon([(bounding_box[0], bounding_box[1]),
                        (bounding_box[2], bounding_box[3]),
                        (bounding_box[4], bounding_box[5]),
                        (bounding_box[6], bounding_box[7])])
        # 计算分割后的块的数量
        num_blocks_x = int(np.ceil(width / block_size))
        num_blocks_y = int(np.ceil(height / block_size))
        # 初始化输出数组
        blocks = np.empty((num_blocks_y, num_blocks_x), dtype=object)
        centers_in_box = np.zeros((num_blocks_y, num_blocks_x), dtype=int)
        blocks_in_box = np.zeros((num_blocks_y, num_blocks_x), dtype=int)

        # 分割图像
        for y in range(num_blocks_y):
            for x in range(num_blocks_x):
                # 计算窗口
                start_x = x * block_size
                end_x = min(start_x + block_size, width)
                start_y = y * block_size
                end_y = min(start_y + block_size, height)

                # 读取块
                window = Window(start_x, start_y, end_x - start_x, end_y - start_y)
                block = src.read(window=window,resampling=Resampling.bilinear)
                blocks[y, x] = block

                # 计算中心点
                center_x = (start_x + end_x) / 2
                center_y = (start_y + end_y) / 2

                # 判断中心点是否位于平行四边形内
                center_point = Point(center_x, center_y)
                if poly.contains(center_point):
                    centers_in_box[y, x] = 1

                # 判断块是否完全位于平行四边形内
                corners = [
                    Point(start_x, start_y),
                    Point(end_x, start_y),
                    Point(start_x, end_y),
                    Point(end_x, end_y)
                ]
                if all(poly.contains(corner) for corner in corners):
                    blocks_in_box[y, x] = 1

            # return np.array(blocks), centers_in_box, blocks_in_box
        save_blocks_based_on_mask(src,block_size,blocks, centers_in_box, out_path, prefix=os.path.splitext(os.path.basename(file_path))[0], suffix='.tif', resampling=Resampling.bilinear)
        return blocks, centers_in_box, blocks_in_box
    # return np.array(blocks,dtype=object).reshape(centers_in_box.shape), centers_in_box, blocks_in_box


# 示例使用
if __name__ == "__main__":
    # TIFF 文件路径
    file_path = 'E:/satLocate/query/5_1007_3452_1907_4252_60.jpg'
    out_path = 'E:/satLocate/query_blocks'
    # 块大小
    block_size = 200

    # 调用函数
    blocks, centers_in_box, blocks_in_box = read_tif_and_split_into_blocks(file_path, block_size,out_path)
    # save_blocks_based_on_mask(file_path,block_size,blocks,centers_in_box,out_path)

    # 打印结果
    print("Blocks shape:", blocks.shape)
    print("Centers in box:\n", centers_in_box)
    print("Blocks in box:\n", blocks_in_box)